
from flask_restful import Resource
from flask import request
from sqlalchemy.orm import relationship
from modules.models import ModelsModel
from modules.dataset import DatasetModel
from common.resp import ajax_success, ajax_error, ajax_data, ajax_page
from core import db
from datetime import datetime

class TrainingModel(db.Model):
    __tablename__ = "sys_training"

    id = db.Column(db.Integer, autoincrement=True, primary_key=True)
    name = db.Column(db.String(50), nullable=False, index=True)
    code = db.Column(db.String(50), nullable=False, index=True)
    ds_id = db.Column(db.Integer, nullable=True)  # 数据集
    model_id = db.Column(db.Integer, nullable=True)  # 模型模型
    last_id = db.Column(db.Integer, nullable=True)  # 上一次训练的ID，用于继续训练
    parameters = db.Column(db.JSON, nullable=False, default={})  # 训练参数，包括学习率、批量大小、训练轮数等
    hardware_config = db.Column(db.JSON, nullable=True)  # 硬件配置，包括 CPU、GPU 和内存等
    status = db.Column(db.String(10), default="pending")  # pending 暂停, training训练中, completed 训练完成, failed 失败, cancelled 取消
    start_time = db.Column(db.DateTime, nullable=True)
    end_time = db.Column(db.DateTime, nullable=True)
    log_path = db.Column(db.String(255), nullable=True)
    tensorboard_path = db.Column(db.String(255), nullable=True)
    process_id = db.Column(db.String(255), nullable=True)  # 训练进程的PID
    last_checkpoint = db.Column(db.String(255), nullable=True)  # 最后一次检查点的路径，用于继续训练
    create_time = db.Column(db.DateTime, nullable=True, default=datetime.now)
    update_time = db.Column(db.DateTime, nullable=True, default=datetime.now, onupdate=datetime.now)

    def to_json(self):
        return {
            'id': self.id,
            'name': self.name,
            'code': self.code,
            'ds_id': self.ds_id,
            'model_id': self.model_id,
            'last_id': self.last_id,
            'parameters': self.parameters,
            'hardware_config': self.hardware_config,
            'status': self.status,
            'start_time': self.start_time.strftime('%Y-%m-%d %H:%M:%S') if self.start_time else None,
            'end_time': self.end_time.strftime('%Y-%m-%d %H:%M:%S') if self.end_time else None,
            'log_path': self.log_path,
            'tensorboard_path': self.tensorboard_path,
            'process_id': self.process_id,
            'last_checkpoint': self.last_checkpoint,
            'create_time': self.create_time.strftime('%Y-%m-%d %H:%M:%S') if self.create_time else None,
            'update_time': self.update_time.strftime('%Y-%m-%d %H:%M:%S') if self.update_time else None,
        }

class TrainingResource(Resource):
    def get(self):
        query = request.args;
        filters = []
        if query.get('name'):
            filters.append(TrainingModel.name.like('%' + query['name'] + '%'))
        if query.get('code'):
            filters.append(TrainingModel.code.like('%' + query['code'] + '%'))
        if query.get('status'):
            filters.append(TrainingModel.status == query['status'])
        if query.get('model_id'):
            filters.append(TrainingModel.model_id == int(query['model_id']))

        paginate = TrainingModel.query.join(DatasetModel, TrainingModel.ds_id == DatasetModel.id).join(ModelsModel, ModelsModel.id == TrainingModel.model_id).filter(
            *filters).paginate(
            page=int(query.get('page', 1)), per_page=int(query.get('limit', 10)), error_out=False)

        return ajax_page([item.to_json() for item in paginate.items], paginate.total)

    def post(self):
        data = request.get_json()
        training = TrainingModel(
            name = data.get('name'),
            code = data.get('code'),
            ds_id = data.get('ds_id', None),
            model_id = data.get('model_id', None),
            last_id = data.get('last_id', None),
            parameters = data.get('parameters', {}),
            hardware_config = data.get('hardware_config', {}),
            status = data.get('status', 'pending')
        )

        db.session.add(training)
        db.session.commit()
        return ajax_success('训练任务创建成功')

    def put(self):
        data = request.get_json()
        if not data.get('id'):
            return ajax_error('参数错误')

        training = TrainingModel.query.get_or_404(data.get('id'))
        if not training:
            return ajax_error('训练任务不存在')

        if training.get('name'):
            training.name = data.get('name', training.name)
        if data.get('code'):
            training.code = data.get('code', training.code)
        if data.get('ds_id'):
            training.ds_id = data.get('ds_id', training.ds_id)
        if data.get('model_id'):
            training.model_id = data.get('model_id', training.model_id)
        if data.get('last_id'):
            training.last_id = data.get('last_id', training.last_id)
        if data.get('parameters'):
            training.parameters = data.get('parameters', training.parameters)
        if data.get('hardware_config'):
            training.hardware_config = data.get('hardware_config', training.hardware_config)
        if data.get('status'):
            training.status = data.get('status', training.status)
        if data.get('start_time'):
            training.start_time = data.get('start_time', training.start_time)
        if data.get('end_time'):
            training.end_time = data.get('end_time', training.end_time)

        db.session.add(training)
        db.session.commit()
        return ajax_success('训练任务更新成功')


class TrainingDetail(Resource):
    def get(self, id):
        if id:
            model = TrainingModel.query.get_or_404(id)
            return ajax_data(model.to_json())
        else:
            return ajax_error('参数错误')

    def delete(self, id):
        train = TrainingModel.query.get_or_404(id)
        if not train:
            return ajax_error('训练任务不存在')

        db.session.delete(train)
        db.session.commit()
        return ajax_success('训练任务删除成功')
